package ocr

import (
	"bytes"
	"encoding/json"
	"errors"
	"image"
	"image/color"
	"io"
	"log"
	"math"
	"net/http"
	"path"
	"path/filepath"
	"sort"
	"strings"

	pd "gitee.com/extrame/tujiang/paddle"
	"github.com/sirupsen/logrus"

	"gocv.io/x/gocv"
)

type PaddleModel struct {
	predictor *pd.Predictor
	inputs    map[string]*pd.Tensor
	outputs   map[string]*pd.Tensor
	inNames   []string
	outNames  []string

	UseGPU             bool   `goblet:"use_gpu" default:"false"`
	DeviceID           int32  `goblet:"gpu_id" default:"0"`
	InitGPUMem         uint64 `goblet:"gpu_mem" default:"1000"`
	NumThreads         int    `goblet:"num_cpu_threads" default:"6"`
	UseMKLDNN          bool   `goblet:"enable_mkldnn" default:"false"`
	UseTensorRT        bool   `goblet:"use_tensorrt" default:"false"`
	UseTrtDynamicShape bool   `goblet:"use_trt_dynamic_shape" default:"false"`
	UseIROptim         bool   `goblet:"ir_optim" default:"true"`
	BatchSize          int    `goblet:"batch_size" default:"1"`
	ModleDir           string `goblet:"model_dir"`

	name string `goblet:"name"`
}

func NewPaddleModel(name string, config *Config) *PaddleModel {
	var model PaddleModel
	model.UseGPU = config.UseGPU
	model.DeviceID = config.DeviceID
	model.InitGPUMem = config.InitGPUMem
	model.NumThreads = config.NumThreads
	model.UseMKLDNN = config.UseMKLDNN
	model.UseTensorRT = config.UseTensorRT
	model.UseTrtDynamicShape = config.UseTrtDynamicShape
	model.UseIROptim = config.UseIROptim
	model.BatchSize = config.BatchSize
	model.ModleDir = config.ModleDir
	model.name = name
	return &model
}

func (model *PaddleModel) LoadModel(modelDir string) {
	config := pd.NewConfig()
	config.DisableGlogInfo()

	config.SetModel(modelDir+"/model", modelDir+"/params")
	if model.UseGPU {
		logrus.WithField("model", model.name).WithField("dir", modelDir).Info("run in gpu mode")
		config.EnableUseGpu(model.InitGPUMem, model.DeviceID)
		if model.UseTensorRT {
			config.EnableTensorRtEngine(1<<30, 16, 3, pd.PrecisionFloat32, false, false)
			if model.UseTrtDynamicShape {
				minInputShape := make(map[string][]int32)
				minInputShape["inputs"] = []int32{int32(model.BatchSize), 3, 100, 100}
				maxInputShape := make(map[string][]int32)
				maxInputShape["inputs"] = []int32{int32(model.BatchSize), 3, 608, 608}
				optInputShape := make(map[string][]int32)
				optInputShape["inputs"] = []int32{int32(model.BatchSize), 3, 224, 224}
				config.SetTRTDynamicShapeInfo(minInputShape, maxInputShape, optInputShape, false)
			}
		}
	} else {
		logrus.WithField("model", model.name).Info("run in cpu mode")
		// config.DisableGpu()
		config.SetCpuMathLibraryNumThreads(model.NumThreads)
	}

	// config.EnableMemoryOptim()
	if model.UseIROptim {
		config.SwitchIrOptim(true)
	}

	// false for zero copy tensor
	// config.SwitchUseFeedFetchOps(false)
	// config.SwitchSpecifyInputNames(true)

	model.predictor = pd.NewPredictor(config)

	model.inNames = model.predictor.GetInputNames()
	model.outNames = model.predictor.GetOutputNames()
	model.inputs = make(map[string]*pd.Tensor)
	model.outputs = make(map[string]*pd.Tensor)

	logrus.WithField("inputs", model.inNames).WithField("outputs", model.outNames).Info("model input/output names")

	for _, n := range model.inNames {
		model.inputs[n] = model.predictor.GetInputHandle(n)
	}
	for _, n := range model.outNames {
		model.outputs[n] = model.predictor.GetOutputHandle(n)
	}

}

type OCRText struct {
	BBox  [][]int `json:"bbox"`
	Text  string  `json:"text"`
	Score float64 `json:"score"`
}

type TextPredictSystem struct {
	detector *DBDetector
	cls      *TextClassifier
	rec      *TextRecognizer
}

func NewTextPredictSystem(config *Config) *TextPredictSystem {
	sys := &TextPredictSystem{
		detector: NewDBDetector(config.DetModelDir, config),
		rec:      NewTextRecognizer(config.RecModelDir, config),
	}
	if config.UseAngleCls {
		sys.cls = NewTextClassifier(config.ClsModelDir, config)
	}
	return sys
}

func (sys *TextPredictSystem) sortBoxes(boxes [][][]int) [][][]int {
	sort.Slice(boxes, func(i, j int) bool {
		if boxes[i][0][1] < boxes[j][0][1] {
			return true
		}
		if boxes[i][0][1] > boxes[j][0][1] {
			return false
		}
		return boxes[i][0][0] < boxes[j][0][0]
	})

	for i := 0; i < len(boxes)-1; i++ {
		if math.Abs(float64(boxes[i+1][0][1]-boxes[i][0][1])) < 10 && boxes[i+1][0][0] < boxes[i][0][0] {
			boxes[i], boxes[i+1] = boxes[i+1], boxes[i]
		}
	}
	return boxes
}

func (sys *TextPredictSystem) getRotateCropImage(img gocv.Mat, box [][]int) gocv.Mat {
	cropW := int(math.Sqrt(math.Pow(float64(box[0][0]-box[1][0]), 2) + math.Pow(float64(box[0][1]-box[1][1]), 2)))
	cropH := int(math.Sqrt(math.Pow(float64(box[0][0]-box[3][0]), 2) + math.Pow(float64(box[0][1]-box[3][1]), 2)))
	ptsstd := make([]image.Point, 4)
	ptsstd[0] = image.Pt(0, 0)
	ptsstd[1] = image.Pt(cropW, 0)
	ptsstd[2] = image.Pt(cropW, cropH)
	ptsstd[3] = image.Pt(0, cropH)

	points := make([]image.Point, 4)
	points[0] = image.Pt(box[0][0], box[0][1])
	points[1] = image.Pt(box[1][0], box[1][1])
	points[2] = image.Pt(box[2][0], box[2][1])
	points[3] = image.Pt(box[3][0], box[3][1])

	M := gocv.GetPerspectiveTransform(gocv.NewPointVectorFromPoints(points), gocv.NewPointVectorFromPoints(ptsstd))
	defer M.Close()
	dstimg := gocv.NewMat()
	gocv.WarpPerspectiveWithParams(img, &dstimg, M, image.Pt(cropW, cropH),
		gocv.InterpolationCubic, gocv.BorderReplicate, color.RGBA{0, 0, 0, 0})

	if float64(dstimg.Rows()) >= float64(dstimg.Cols())*1.5 {
		srcCopy := gocv.NewMat()
		gocv.Transpose(dstimg, &srcCopy)
		defer dstimg.Close()
		gocv.Flip(srcCopy, &srcCopy, 0)
		return srcCopy
	}
	return dstimg
}

func (sys *TextPredictSystem) Run(img gocv.Mat) []OCRText {
	srcimg := gocv.NewMat()
	img.CopyTo(&srcimg)
	boxes := sys.detector.Run(img)
	if len(boxes) == 0 {
		return nil
	}

	boxes = sys.sortBoxes(boxes)
	cropimages := make([]gocv.Mat, len(boxes))
	for i := 0; i < len(boxes); i++ {
		tmpbox := make([][]int, len(boxes[i]))
		for j := 0; j < len(tmpbox); j++ {
			tmpbox[j] = make([]int, len(boxes[i][j]))
			copy(tmpbox[j], boxes[i][j])
		}
		cropimg := sys.getRotateCropImage(srcimg, tmpbox)
		cropimages[i] = cropimg
	}
	if sys.cls != nil {
		cropimages = sys.cls.Run(cropimages)
	}
	recResult := sys.rec.Run(cropimages, boxes)
	return recResult
}

type OCRSystem struct {
	// args *Config
	tps *TextPredictSystem
}

func NewOCRSystem(config *Config) *OCRSystem {
	logrus.Info("NewOCRSystem")

	config.Check()

	return &OCRSystem{
		// args: config,
		tps: NewTextPredictSystem(config),
	}
}

func (ocr *OCRSystem) StartServer(port string) {
	http.HandleFunc("/ocr", ocr.predictHandler)
	log.Println("OCR Server has been started on port :", port)
	err := http.ListenAndServe(":"+port, nil)
	if err != nil {
		log.Panicf("http error! error: %v\n", err)
	}
}

func (ocr *OCRSystem) predictHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.Write([]byte(errors.New("post method only").Error()))
		return
	}
	r.ParseMultipartForm(32 << 20)
	var buf bytes.Buffer
	file, header, err := r.FormFile("image")
	if err != nil {
		w.Write([]byte(err.Error()))
		return
	}
	defer file.Close()
	ext := strings.ToLower(path.Ext(header.Filename))
	if ext != ".jpg" && ext != ".png" {
		w.Write([]byte(errors.New("only support image endswith jpg/png").Error()))
		return
	}

	io.Copy(&buf, file)
	img, err2 := gocv.IMDecode(buf.Bytes(), gocv.IMReadColor)
	if err2 != nil {
		w.Write([]byte(err2.Error()))
		return
	}
	result := ocr.PredictOneImage(img)
	if output, err3 := json.Marshal(result); err3 != nil {
		w.Write([]byte(err3.Error()))
	} else {
		w.Write(output)
	}
}

func (ocr *OCRSystem) PredictOneImage(img gocv.Mat) []OCRText {
	return ocr.tps.Run(img)
}

func (ocr *OCRSystem) PredictDirImages(dirname string) map[string][]OCRText {
	if dirname == "" {
		return nil
	}

	imgs, _ := filepath.Glob(dirname + "/*.jpg")
	tmpimgs, _ := filepath.Glob(dirname + "/*.png")
	imgs = append(imgs, tmpimgs...)
	results := make(map[string][]OCRText, len(imgs))
	for i := 0; i < len(imgs); i++ {
		imgname := imgs[i]
		img, err := ReadImage(imgname)
		if err != nil {
			log.Printf("read image %v failed! err: %v\n", imgname, err)
			continue
		}
		res := ocr.PredictOneImage(img)
		results[imgname] = res
	}
	return results
}
